#include "psrfold.h"
#include "data.h"

#define PI 3.14159265358979323846

const double dm = 2.64476;
const double period = 0.00575730363767324;

long get_period_size(float bw)
{
    return period * bw * 1e6;
}

void apply_chirp(fftw_complex *fdata,float bw,float fc,long size)
{
    const double dm_dispersion = 2.41e-4;
    double dispersion_per_MHz = 1e6 * dm / dm_dispersion;
    // size = size / 2;
    fftw_complex *phasors = malloc(sizeof(fftw_complex) * size);
    double binwidth = -bw / size;
    double coeff = 2 * PI * dispersion_per_MHz / (fc * fc);
    for(int k = 0; k < size; ++k){
        double freq = k * binwidth + 0.5 * bw;
        phasors[k] = cexp(I * coeff * freq * freq / (fc+freq));
        // phasors[k] = exp(i*phases[k]);
        if(k == 0) phasors[k] = 0; 
        fdata[k] = phasors[k] * fdata[k];
    }
    free(phasors);
}

void fft(fftw_complex *in,fftw_complex *out,long size)
{
    fftw_plan_with_nthreads(4);
    fftw_plan plan = fftw_plan_dft_1d(size,in,out,FFTW_FORWARD,FFTW_ESTIMATE);
    if(plan == NULL) {
        printf("fft:Create fftw_plan error!");
        exit(1);
    }
    fftw_execute(plan);
    fftw_destroy_plan(plan);
    fftw_cleanup_threads();
}

void ifft(fftw_complex *in,double *out,long size)
{
    fftw_plan_with_nthreads(4);
    fftw_plan plan = fftw_plan_dft_c2r_1d(size,in,out,FFTW_ESTIMATE);
    if(plan == NULL) {
        printf("ifft:Create fftw_plan error!");
        exit(1);
    }
    fftw_execute(plan);
    fftw_destroy_plan(plan);
    fftw_cleanup_threads();
}

void ifft_c2c(fftw_complex *in,fftw_complex *out,long size)
{
    fftw_plan_with_nthreads(4);
    fftw_plan plan = fftw_plan_dft_1d(size,in,out,FFTW_BACKWARD,FFTW_ESTIMATE);
    if(plan == NULL) {
        printf("ifft:Create fftw_plan error!");
        exit(1);
    }
    fftw_execute(plan);
    fftw_destroy_plan(plan);
    fftw_cleanup_threads();
}

void amplitude(double *amp1,double *amp2,double *out,long size)
{
    for(int i = 0;i < size;++i){
        out[i] = sqrt(amp1[i]*amp1[i] + amp2[i]*amp2[i]);
    }
}

void compute_amplitude(fftw_complex *fdata,double * out,long size)
{
    for(int i = 0;i < size; ++i){
        out[i] = sqrt(creal(fdata[i])*creal(fdata[i]) + cimag(fdata[i])*cimag(fdata[i]));
    }
}

void amplitude1(double *amp1,double *amp2,double *out,long size)
{
    for(int i = 0;i < size;++i){
        out[i] = amp1[i];
    }
}

void coherent_dedispersion1(fftw_complex *pol1,fftw_complex *pol2,double *out,float bw,float fc,long size)
{
    fftw_complex *out1 = malloc(sizeof(fftw_complex) * size);
    fftw_complex *out2 = malloc(sizeof(fftw_complex) * size);
    // fftw_complex *res11 = malloc(sizeof(fftw_complex) * size);
    // fftw_complex *res22 = malloc(sizeof(fftw_complex) * size);
    double * res1 = malloc(sizeof(double) * size);
    double * res2 = malloc(sizeof(double) * size);

    fft(pol1,out1,size);
    fft(pol1,out2,size);

    apply_chirp(out1,bw,fc,size);
    apply_chirp(out2,bw,fc,size);

    ifft(out1,res1,size);
    ifft(out2,res2,size);

    // compute_amplitude(res11, res1, size);
    // compute_amplitude(res22, res2, size);

    amplitude1(res1,res2,out,size);

    fftw_free(out1);
    fftw_free(out2);
    // fftw_free(res11);
    // fftw_free(res22);
    free(res1);
    free(res2);
}

void coherent_dedispersion(fftw_complex *pol1,fftw_complex *pol2,double *out,float bw,float fc,long size,int index,int nfft)
{
    fftw_complex *out1 = malloc(sizeof(fftw_complex) * size);
    fftw_complex *out2 = malloc(sizeof(fftw_complex) * size);
    double * res1 = malloc(sizeof(double) * size);
    double * res2 = malloc(sizeof(double) * size);

    fft(pol1,out1,size);
    fft(pol1,out2,size);

    long start = 0;
    long valid_size = 0;

    // extract valid freq
    extract_freq(out1,size,&valid_size,&start,index,nfft);
    // extract_freq(out2,size,&valid_size,&start,index,nfft);

    // printf("valid_size = %ld\n",valid_size);
   
    apply_chirp(out1+start,bw,fc,valid_size);
    apply_chirp(out2+start,bw,fc,valid_size);

    ifft(out1+start,res1,valid_size);
    ifft(out2+start,res2,valid_size);

    amplitude(res1,res2,out,valid_size);

    fftw_free(out1);
    fftw_free(out2);
    free(res1);
    free(res2);
}

void combine_freq(fftw_complex * full_freq,fftw_complex * sub_time,int index,int nfft,long size)
{
    fftw_complex * sub_freq = malloc(sizeof(fftw_complex) * size);
    fft(sub_time,sub_freq,size);

    long start = 0;
    long valid_size = 0;
    long combine_start = 0;

    for(int i = 0;i < index;++i){
        if(i == 0){
            combine_start += size / 4;
        }else{
            combine_start += size / 2;
        }
    }

    // extract valid freq
    extract_freq(sub_freq,size,&valid_size,&start,index,nfft);

    // printf("index = %d,start = %ld,end = %ld\n",index,combine_start,combine_start+valid_size);

    for(int i = 0;i < valid_size;++i){
        full_freq[i+combine_start] = sub_freq[i+start];
    }
}

void combine_freq1(fftw_complex * full_freq,fftw_complex * sub_time,int index,int nfft,long size)
{
    fftw_complex * sub_freq = malloc(sizeof(fftw_complex) * size);
    fft(sub_time,sub_freq,size);

    long combine_start = index * size;

    for(int i = 0;i < size;++i){
        full_freq[i+combine_start] = sub_freq[i];
    }
}

void extract_subband(fftw_complex *data,fftw_complex *subband,long size,int nfft,int index)
{
    for(int i = 0,j = 0;i < size;){
        subband[j] = data[i+index];
        j += 1;
        i += nfft;
    }
}

void shift_freq(fftw_complex *freq,long size)
{
    fftw_complex temp;
    for(int i = 0;i < size / 2;++i){
        temp = freq[i];
        freq[i] = freq[i+size/2];
        freq[i+size/2] = temp;
    }
}

void flip_freq(fftw_complex *freq,long size)
{
    fftw_complex temp;
    for(int i = 0;i < size / 2;++i){
        temp = freq[i];
        freq[i] = freq[size-i-1];
        freq[size-i-1]= temp;
    }
}

// void shift_freq(fftw_complex *freq,long size)
// {
//     double real,imag;
//     for(int i = 0;i < size / 2;++i){
//         real = freq[i][0];
//         imag = freq[i][1];
//         freq[i][0] = freq[i+size/2][0];
//         freq[i][1] = freq[i+size/2][1];
//         freq[i+size/2][0] = real;
//         freq[i+size/2][1] = imag;
//     }
// }

void extract_freq(fftw_complex *freq,long size,long *valid_size,long *start,int index,int nfft)
{
    if(index % 2 == 0){
        shift_freq(freq,size);
    }
    // flip_freq(freq,size);
    // for 0 and nfft / 2 subband, only extract size / 4
    if(index == 0){
        *start = size / 2;
        *valid_size = size / 4;
    }else if(index == (nfft / 2)){
        *start = size / 4;
        *valid_size = size / 4;
    }else{
        *start = size / 4;
        *valid_size = size / 2;
    }
}

void write_data_c(const char *filename,fftw_complex *data,long size)
{
    FILE *pf = fopen(filename,"w");
    if(pf == NULL){
        perror("write_data:fopen()");
        exit(1);
    }

    for(int i = 0;i < size;++i){
        fprintf(pf,"%f\n",creal(data[i]));
        fprintf(pf,"%f\n",cimag(data[i]));
    }

    printf("Write successfully!\n");
    fclose(pf);
}

void dedispersion_fold(fftw_complex *subband1,fftw_complex *subband2,double *pdata,double *pnum,long *location,int nfft,long psize,long size)
{
    int num = nfft / 2 + 1;
    float bandwidth = 400.0 / (num-1);
    long subband_size = size / nfft;
    long real_size = subband_size * num;
    // printf("num = %d, bandwidth = %f,real_size = %ld\n",num,bandwidth,real_size);
    float bw = 0.0;
    float fc = 1582.0;
    
    fftw_complex * pol_sub_1 = malloc(sizeof(fftw_complex) * subband_size);
    fftw_complex * pol_sub_2 = malloc(sizeof(fftw_complex) * subband_size);
    long n = subband_size * (num-1) / 2;
    fftw_complex * full_freq_1 = malloc(sizeof(fftw_complex) * n);
    fftw_complex * full_freq_2 = malloc(sizeof(fftw_complex) * n);
    double * pol_out = malloc(sizeof(double) * subband_size);
    double * out = malloc(sizeof(double) * n);
    long block_size = 0;
    long current_psize = 0;

    // printf("n = %ld,block_size = %ld\n",n,size / 2);

    for(int i = 0;i < num;++i){
        // so guly, bandwidth not equal
        fc -= bw / 2;
        if(i == 0 || i == num-1){
            bw = bandwidth / 2;
            block_size = subband_size / 2;
            current_psize = psize / 2;
        }else{
            bw = bandwidth;
            block_size = subband_size;
            current_psize = psize;
        }
        // current_psize = get_period_size(bw);
        fc -= bw / 2;
        printf("current_psize = %ld,psize = %ld,fc = %f,bw = %f,block_size = %ld\n",current_psize,psize,fc,bw,block_size);
        // printf("begin extract_subband\n");
        extract_subband(subband1,pol_sub_1,size,nfft,i);
        extract_subband(subband2,pol_sub_2,size,nfft,i);

        combine_freq(full_freq_1,pol_sub_1,i,nfft,subband_size);
        combine_freq(full_freq_2,pol_sub_2,i,nfft,subband_size);

        // printf("begin coherent_dedispersion\n");
        coherent_dedispersion(pol_sub_1,pol_sub_2,pol_out,bw,fc,subband_size,i,nfft);
        // printf("begin fold_data\n");

        fold_data(pdata+i*psize,pnum+i*psize,pol_out,current_psize,block_size/2,&location[i]);
    }
    // coherent_dedispersion1(full_freq_1,full_freq_2,out,400.0,1382.0,n);

    // ifft(full_freq_1,out,n);
    // printf("before location = %ld\n",location[0]);
    // fold_data(pdata,pnum,out,psize,n,&location[0]);
    // printf("after location = %ld\n",location[0]);
    // compute_amplitude(full_freq_1, out, n);
    write_data_c("freq.txt",full_freq_1,n);
    fftw_free(pol_sub_1);
    fftw_free(pol_sub_2);
    fftw_free(full_freq_1);
    fftw_free(full_freq_2);
    // free(out);

    free(pol_out);
}

void fold_data(double *pdata,double *pnum,double *bdata,long psize,long blocksize,long *loc)
{
    // printf("location = %ld\n",*loc);
    long cur = *loc;
    for(int i = 0;i < blocksize;++i){
        if(cur >= psize){
            cur = 0;
        }
        pnum[cur]++;
        pdata[cur] = (pnum[cur] - 1) * pdata[cur] / pnum[cur] + bdata[i] / pnum[cur];
        ++cur;
    }
    *loc = cur;
}

void integral_data(double *data,double *out,long size)
{
    int bins = size / nbin;
    int index = 0;
    double sum;
    long i = 0;
    for(;i < bins * (nbin-1);){
        sum = 0;
        long j = 0;
        for(;j < bins && i < bins * (nbin-1);++j){
            sum += data[i];
            ++i;
        }
        out[index] = sum / j;
        ++index;
    }
    sum = 0;
    for(;i < size;++i){    
        sum += data[i];
    }
    out[nbin-1] = sum / (size - nbin * bins + bins);
}

void offset_data(double *data,double *out,double phi)
{
    int offset = phi * nbin;
    // double *new_data = malloc(sizeof(double) * nbin);
    // if(new_data == NULL){
    //     perror("offset_data:malloc()");
    //     exit(1);
    // }
    int i = 0;
    for(; i < offset; ++i){
        out[i] = data[nbin - offset + i];
    }
    for(; i < nbin; ++i){
        out[i] = data[i - offset];
    }
}
